#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: runyuanye@163.com
"""
import torch
import numpy as np


def euclidean_squared_distance(input):
    """Computes euclidean squared distance.

    Args:
        input (torch.Tensor): input, feature matrix.

    Returns:
        torch.Tensor: distance matrix.
    """
    m = n = input.size(0)
    sq_input = torch.pow(input, 2).sum(dim=1, keepdim=True)
    distmat = sq_input.expand(m, n) + \
              sq_input.expand(n, m).t()
    distmat.addmm_(1, -2, input, input.t())
    distmat[distmat < 0] = 0
    return distmat


def euclidean_squared_distance2(input1, input2):
    """Computes euclidean squared distance.

    Args:
        input1 (torch.Tensor): 2-D feature matrix.
        input2 (torch.Tensor): 2-D feature matrix.

    Returns:
        torch.Tensor: distance matrix.
    """
    m, n = input1.size(0), input2.size(0)
    distmat = torch.pow(input1, 2).sum(dim=1, keepdim=True).expand(m, n) + \
              torch.pow(input2, 2).sum(dim=1, keepdim=True).expand(n, m).t()
    distmat.addmm_(1, -2, input1, input2.t())
    return distmat


def cosine_distance(input1):
    """Computes cosine distance.

    Args:
        input1 (torch.Tensor): 2-D feature matrix.

    Returns:
        torch.Tensor: distance matrix.
    """
    distmat = 1 - torch.mm(input1, input1.t())
    distmat.clamp_(0, 1)
    return distmat


def cosine_distance2(input1, input2):
    """Computes cosine distance.

    Args:
        input1 (torch.Tensor): 2-D feature matrix.
        input2 (torch.Tensor): 2-D feature matrix.

    Returns:
        torch.Tensor: distance matrix.
    """
    distmat = 1 - torch.mm(input1, input2.t())
    distmat.clamp_(0, 1)
    return distmat


def cosine_distance_ex(input1):
    n, dim = input1.size(0), input1.size(1)
    max_n = 1024
    process_count = (n * n + max_n * max_n - 1) // (max_n * max_n)
    if process_count < 1:
        process_count = 1
    if process_count == 1:
        distmat = cosine_distance(input1)
        return distmat.sum(axis=1)
    process_n = (n + process_count - 1) // process_count
    # print(n, process_n, process_count)
    offsets = [i * process_n for i in range(process_count)]
    offsets.append(n)
    dist_sums = []
    for i in range(process_count):
        input2 = input1[offsets[i]:offsets[i+1], ...]
        distmat = cosine_distance2(input1, input2)
        distmat = distmat.t()
        dist_sum = distmat.sum(axis=1)
        dist_sums.append(dist_sum)
    return torch.cat(dist_sums, 0)
